#include "Basis.h"
#include "Drv.h"
#include "Cpu.h"
#include "KeyMouse.h"
#include "Common.h"
#include "Notification.h"
#include "Hook.h"
#include "Network.h"

NTSTATUS DriverEntry(__in PDRIVER_OBJECT DriverObject, __in PUNICODE_STRING RegistryPath)
{
	UNICODE_STRING		DeviceName, Win32Device;
	PDEVICE_OBJECT		DeviceObject = NULL;
	NTSTATUS			status;
	ULONG				i;

	UNREFERENCED_PARAMETER(RegistryPath);

	RtlInitUnicodeString(&DeviceName, AST_DEVICE_NAME);
	RtlInitUnicodeString(&Win32Device, AST_DOS_DEVICE_NAME);

	// set handlers
	for (i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; i++)
	{
		DriverObject->MajorFunction[i] = DefaultHandler;
	}
	DriverObject->MajorFunction[IRP_MJ_CREATE] = CreateCloseHandler;
	DriverObject->MajorFunction[IRP_MJ_CLOSE] = CreateCloseHandler;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DeviceControlHandler;
	DriverObject->DriverUnload = DriverUnload;

	// create device
	status = IoCreateDevice(DriverObject, 0, &DeviceName, FILE_DEVICE_AST, FILE_DEVICE_SECURE_OPEN, FALSE, &DeviceObject);
	if (!NT_SUCCESS(status)) return status;

	// create dos device
	status = IoCreateSymbolicLink(&Win32Device, &DeviceName);
	if (!NT_SUCCESS(status))
	{
		IoDeleteDevice(DeviceObject);
		return status;
	}

	// initialization
	InitializeCpu();
	InitializeNotification();
	InitializeKeyMouse();
	InitializeHook();
	InitializeNetwork();

	DeviceObject->Flags |= DO_DIRECT_IO;
	DeviceObject->AlignmentRequirement = FILE_WORD_ALIGNMENT;
	DeviceObject->Flags &= ~DO_DEVICE_INITIALIZING;

	return STATUS_SUCCESS;
}

VOID DriverUnload(__in PDRIVER_OBJECT DriverObject)
{
	UNICODE_STRING Win32Device;

	// dispose
	DisposeCpu();
	DisposeNotification();
	DisposeKeyMouse();
	DisposeHook();
	DisposeNetwork();

	// unload driver
	RtlInitUnicodeString(&Win32Device, AST_DOS_DEVICE_NAME);
	IoDeleteSymbolicLink(&Win32Device);
	IoDeleteDevice(DriverObject->DeviceObject);
}

NTSTATUS DefaultHandler(__in PDEVICE_OBJECT DeviceObject, __in PIRP Irp)
{
	UNREFERENCED_PARAMETER(DeviceObject);

	Irp->IoStatus.Status = STATUS_NOT_SUPPORTED;
	Irp->IoStatus.Information = 0;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return Irp->IoStatus.Status;
}

NTSTATUS CreateCloseHandler(__in PDEVICE_OBJECT DeviceObject, __in PIRP Irp)
{
	UNREFERENCED_PARAMETER(DeviceObject);

	Irp->IoStatus.Status = STATUS_SUCCESS;
	Irp->IoStatus.Information = 0;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS DeviceControlHandler(__in PDEVICE_OBJECT DeviceObject, __in PIRP Irp)
{
	PIO_STACK_LOCATION irpStack = IoGetCurrentIrpStackLocation(Irp);
	LONG ioControlCode = irpStack->Parameters.DeviceIoControl.IoControlCode;
	
	Irp->IoStatus.Information = 0;
	
	BOOLEAN bFalse = FALSE;
	do
	{
		// processing
		if (ProcessIoCtlForKeyMouse(ioControlCode, DeviceObject, Irp, irpStack)) break;
		if (ProcessIoCtlForCommon(ioControlCode, DeviceObject, Irp, irpStack)) break;
		if (ProcessIoCtlForHook(ioControlCode, DeviceObject, Irp, irpStack)) break;
		if (ProcessIoCtlForNetwork(ioControlCode, DeviceObject, Irp, irpStack)) break;

		// unknown io control code
		Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
		KdPrint(("\nunknown IRP_MJ_DEVICE_CONTROL\n"));

	} while (bFalse);
	
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return Irp->IoStatus.Status;
}

